#include "dxut.h"
#include "dx_util.h"
#include "scene_renderer.h"

// note: this has to exactly match the order in the hlsl shader file
#define IDSRV_TXOBJ		0
#define IDSRV_TXSM		1
#define IDSRV_TXNOISE	2

struct CB_MAT_OBJ
{
	D3DXVECTOR4		m_MatAmbient;
	D3DXVECTOR4		m_MatDiffuse; // alpha value: whether texture is used (yes: 1; no: -1)
	D3DXVECTOR4		m_MatSpecular; // alpha value: Specular Exponent
};

struct CB_PER_FRAME
{
	D3DXMATRIX	m_mViewProj;
	D3DXMATRIX  m_mViewProjLight;

	D3DXVECTOR3 m_vEye;
	float		m_dummy;

	D3DXVECTOR4 m_AmbientLightColor;

	D3DXVECTOR4 m_PointLightColor;
	D3DXVECTOR4 m_PointLightPos;

	D3DXVECTOR3 m_vLookAt;
	float		m_fvFOV;
	
    D3DXMATRIX  m_mInvViewProj;
	D3DXMATRIX  m_mInvViewProjLight;
	
	int			m_iTechShading;
	int			m_iTechSM;
    int         m_todo;
    int         m_todo1;
};
CSceneRenderer::CSceneRenderer()
{}

CSceneRenderer::~CSceneRenderer()
{}

void CSceneRenderer::CreateShaderAndLayout(ID3D11Device *pd3dDevice)
{
    // Compile the shaders
    DWORD dwShaderFlags = D3D10_SHADER_ENABLE_STRICTNESS;
#if defined( DEBUG ) || defined( _DEBUG )
    // Set the D3D10_SHADER_DEBUG flag to embed debug information in the shaders.
    // Setting this flag improves the shader debugging experience, but still allows 
    // the shaders to be optimized and to run exactly the way they will run in 
    // the release configuration of this program.
    dwShaderFlags |= D3D10_SHADER_DEBUG;
#endif

	CreateVertexShader(pd3dDevice, L"shaders\\renderscene.hlsl", "VS_RenderScene", dwShaderFlags, &m_pVSScene);
	CreatePixelShader(pd3dDevice, L"shaders\\renderscene.hlsl", "PS_SceneShading", dwShaderFlags, &m_pPSScene);

	//////////////////////////////////////////////////////////////////////////////////////////////////////
	// Create a layout for the hair vertex data
	ID3DBlob* pVertexShaderBuffer = NULL;
	ID3DBlob* pErrorBuffer = NULL;

	//////////////////////////////////////////////////////////////////////////////////////////////////////
	// Create a layout for the scene object data
	if( FAILED( D3DX11CompileFromFile( L"shaders\\renderScene.hlsl", NULL, NULL, "VS_RenderScene", "vs_5_0", dwShaderFlags, 0, NULL,
		&pVertexShaderBuffer, &pErrorBuffer, NULL ) ))
		MessageBoxA(NULL, (LPCSTR)pErrorBuffer->GetBufferPointer(), "VS Compile Error",0);
	SAFE_RELEASE(pErrorBuffer);

	const D3D11_INPUT_ELEMENT_DESC layout_mesh[] =
	{
		{ "POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0,  0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
		{ "NORMAL",   0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 12, D3D11_INPUT_PER_VERTEX_DATA, 0 },
		{ "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT,    0, 24, D3D11_INPUT_PER_VERTEX_DATA, 0 },
	};

	if(FAILED( pd3dDevice->CreateInputLayout( layout_mesh, ARRAYSIZE( layout_mesh ), pVertexShaderBuffer->GetBufferPointer(),
		pVertexShaderBuffer->GetBufferSize(), &m_pLayoutScene )))
		::MessageBoxA(0, "Fail to create object layout", "d3d error", 0);
	SAFE_RELEASE( pVertexShaderBuffer );
}

void CSceneRenderer::CreateConstantBuffers(ID3D11Device* pd3dDevice)
{
    HRESULT hr;
	///////////////////////////////////////////////////////////////////////////////////////////////
    // Create constant buffers
    D3D11_BUFFER_DESC cbDesc;
    ZeroMemory( &cbDesc, sizeof(cbDesc) );
    cbDesc.Usage = D3D11_USAGE_DYNAMIC;
    cbDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    cbDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;

    cbDesc.ByteWidth = sizeof( CB_PER_FRAME );
    hr = pd3dDevice->CreateBuffer( &cbDesc, NULL, &m_pcbPerFrame ) ;
    if (FAILED(hr))
	    ::MessageBoxA(0, "Fail to create constant buffer for frame", "d3d error", 0);

    cbDesc.ByteWidth = sizeof( CB_MAT_OBJ );
    hr = pd3dDevice->CreateBuffer( &cbDesc, NULL, &m_pcbMatObj ) ;
    if (FAILED(hr))
	    ::MessageBoxA(0, "Fail to create constant buffer for material", "d3d error", 0);
}

void CSceneRenderer::CreateSamplerStates( ID3D11Device* pd3dDevice )
{
    HRESULT hr;
	//////////////////////////////////////////////////////////////////////////////////////////
	// Create sampler state objects
	D3D11_SAMPLER_DESC samDesc;
	ZeroMemory( &samDesc, sizeof(samDesc) );
	samDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR ;//D3D11_FILTER_MIN_MAG_MIP_LINEAR;
	samDesc.AddressU = samDesc.AddressV = samDesc.AddressW = D3D11_TEXTURE_ADDRESS_WRAP;
	samDesc.MaxAnisotropy = 16;
	samDesc.ComparisonFunc = D3D11_COMPARISON_NEVER;
	samDesc.MinLOD = 0;
	samDesc.MaxLOD = D3D11_FLOAT32_MAX;
	hr = pd3dDevice->CreateSamplerState( &samDesc, &m_pSamplerStateLinearWrap ) ;
	if(FAILED(hr))
		::MessageBoxA(0, "Fail to create linear wrap sampler state", "D3D Error", 0);

	samDesc.AddressU = samDesc.AddressV = samDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
	hr = pd3dDevice->CreateSamplerState( &samDesc, &m_pSamplerStateLinearClamp ) ;
	if(FAILED(hr))
		::MessageBoxA(0, "Fail to create linear clamp sampler state", "D3D Error", 0);

	samDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_POINT;
	hr = pd3dDevice->CreateSamplerState( &samDesc, &m_pSamplerStatePointClamp ) ;
	if(FAILED(hr))
		::MessageBoxA(0, "Fail to create point sampler state", "D3D Error", 0);

	samDesc.Filter = D3D11_FILTER_COMPARISON_MIN_MAG_MIP_LINEAR;//D3D11_FILTER_COMPARISON_MIN_MAG_LINEAR_MIP_POINT;
	samDesc.ComparisonFunc = D3D11_COMPARISON_LESS;
	samDesc.BorderColor[0] = samDesc.BorderColor[1] = 
		samDesc.BorderColor[2] = samDesc.BorderColor[3] = 1;
	samDesc.AddressU = D3D11_TEXTURE_ADDRESS_BORDER;
	samDesc.AddressV = D3D11_TEXTURE_ADDRESS_BORDER;
	samDesc.AddressW = D3D11_TEXTURE_ADDRESS_BORDER;
	if(FAILED(pd3dDevice->CreateSamplerState( &samDesc, &m_pSamplerStateCmpLess )))
		::MessageBoxA(0, "Fail to create compare sampler state", "D3D error", 0);
}

void CSceneRenderer::OnCreateDevice(ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dContext)
{
    CDX11Mesh* pMesh = new CDX11Mesh();
    pMesh->CreateMeshFromOBJFile(pd3dDevice, L"..\\media\\objs\\elephant\\elephant.obj");
    m_SceneObjects.push_back(pMesh);
    // once models are ready, compute bounding box
    CalcBoundBox();

    CreateSamplerStates(pd3dDevice);
    CreateShaderAndLayout(pd3dDevice);
    CreateConstantBuffers(pd3dDevice);
}

void CSceneRenderer::SetSamplerStates( ID3D11DeviceContext* pd3dContext )
{
	// Set texture sampler states
	ID3D11SamplerState* pSS[4];
	pSS[0] = m_pSamplerStateLinearWrap;
	pSS[1] = m_pSamplerStateLinearClamp;
	pSS[2] = m_pSamplerStatePointClamp;
	pSS[3] = m_pSamplerStateCmpLess;

	pd3dContext->PSSetSamplers(0, 4, pSS);
	pd3dContext->VSSetSamplers(0, 4, pSS);
	pd3dContext->HSSetSamplers(0, 4, pSS);
	pd3dContext->DSSetSamplers(0, 4, pSS);
	pd3dContext->GSSetSamplers(0, 4, pSS);
}

void CSceneRenderer::CalcBoundBox()
{
    if (m_SceneObjects.empty())
        return;

    m_BBMax = m_SceneObjects[0]->GetBBMax();
    m_BBMin = m_SceneObjects[0]->GetBBMin();

    for (UINT i = 1; i < m_SceneObjects.size(); i++)
    {
        m_BBMax = D3DXVECTOR3(  max(m_BBMax.x, m_SceneObjects[i]->GetBBMax().x), 
                                max(m_BBMax.y, m_SceneObjects[i]->GetBBMax().y), 
                                max(m_BBMax.z, m_SceneObjects[i]->GetBBMax().z));
        m_BBMin = D3DXVECTOR3(  min(m_BBMin.x, m_SceneObjects[i]->GetBBMin().x),
                                min(m_BBMin.y, m_SceneObjects[i]->GetBBMin().y),
                                min(m_BBMin.z, m_SceneObjects[i]->GetBBMin().z));
    }

}

void CSceneRenderer::SetConstantBuffer(ID3D11DeviceContext* pd3dContext, FRAME_PARAMS& frameParams)
{
	HRESULT hr;

    D3DXVECTOR3 sceneCenter = (m_BBMax + m_BBMin)/2;
    D3DXVECTOR3 sceneVec = m_BBMax - m_BBMin;
    float sceneRadius = D3DXVec3Length(&sceneVec)/2;
	/////////////////////////////////////////////////////////////////////////////////////
	// Set up light camera parameters: //TBD consider Orthogonal projection instead
	D3DXMATRIX mViewLight, mProjLight, mViewProjLight;

	D3DXVECTOR3 vLightPosition = frameParams.m_vPointLightPos;
	D3DXVECTOR3 vLightAt = sceneCenter;
	D3DXVECTOR3 vUp = D3DXVECTOR3(0,1,0);
	D3DXMatrixLookAtLH(&mViewLight, &vLightPosition, &vLightAt, &vUp);

	D3DXVECTOR3 vLightToObject = vLightPosition - vLightAt;
	float dis = D3DXVec3Length(&vLightToObject);

	float min_dis = max(0.1f, dis - sceneRadius);
	float max_dis = dis + sceneRadius;

	float halfAngle = asin(sceneRadius/dis);
	float vfov = 2*halfAngle;

	float ratio = 1;

	D3DXMatrixPerspectiveFovLH(&mProjLight, vfov, ratio, min_dis, max_dis );
	mViewProjLight = mViewLight * mProjLight;

	/////////////////////////////////////////////////////////////////////////////////////////
	// constant buffer per frame
	D3D11_MAPPED_SUBRESOURCE MappedResource;
	hr = pd3dContext->Map( m_pcbPerFrame, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) ;
	if (FAILED(hr))
		::MessageBoxA(0, "Fail to map constant buffer", "d3d error", 0);

	CB_PER_FRAME* pcbPerFrame = ( CB_PER_FRAME* )MappedResource.pData;

	// Light camera parameters
	D3DXMatrixTranspose( &pcbPerFrame->m_mViewProjLight, &mViewProjLight);
	D3DXMATRIX mInvViewProjLight;
	D3DXMatrixInverse(&mInvViewProjLight, 0, &mViewProjLight);
	D3DXMatrixTranspose(&pcbPerFrame->m_mInvViewProjLight, &mInvViewProjLight);
    
	// scene light
	pcbPerFrame->m_PointLightPos = frameParams.m_vPointLightPos;
	pcbPerFrame->m_AmbientLightColor = frameParams.m_vAmbientLightColor;
	pcbPerFrame->m_PointLightColor = frameParams.m_vPointLightColor;

    // camera parameters
	D3DXMATRIX mViewProj = frameParams.m_mView * frameParams.m_mProj;
	D3DXMatrixTranspose(&pcbPerFrame->m_mViewProj, &mViewProj);

	D3DXMATRIX mInvViewProj;
	D3DXMatrixInverse(&mInvViewProj, 0, &mViewProj);
	D3DXMatrixTranspose(&pcbPerFrame->m_mInvViewProj, &mInvViewProj);

	pcbPerFrame->m_vEye = frameParams.m_vEye;
	pcbPerFrame->m_vLookAt = frameParams.m_vLookAt;
	pcbPerFrame->m_fvFOV = frameParams.m_fvFOV;

    pcbPerFrame->m_iTechShading = 2;
    
	pd3dContext->Unmap( m_pcbPerFrame, 0 );

	pd3dContext->VSSetConstantBuffers(0, 1, &m_pcbPerFrame);
	pd3dContext->HSSetConstantBuffers(0, 1, &m_pcbPerFrame);
	pd3dContext->DSSetConstantBuffers(0, 1, &m_pcbPerFrame);
	pd3dContext->GSSetConstantBuffers(0, 1, &m_pcbPerFrame);
	pd3dContext->PSSetConstantBuffers(0, 1, &m_pcbPerFrame);
}

void CSceneRenderer::RenderDX11Mesh(ID3D11DeviceContext* pd3dContext, CDX11Mesh* pMesh)
{
    HRESULT hr;

	////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Render scene 
    UINT strides = sizeof( CDX11Mesh::Vertex );
	UINT offsets = 0;
	pd3dContext->IASetInputLayout( m_pLayoutScene );
    ID3D11Buffer* pVB = pMesh->GetVertexBuffer();
    pd3dContext->IASetVertexBuffers( 0, 1, &(pVB), &strides, &offsets );
    pd3dContext->IASetIndexBuffer( pMesh->GetIndexBuffer(), DXGI_FORMAT_R32_UINT, 0 );
	pd3dContext->IASetPrimitiveTopology( D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST );

	// set shader
	pd3dContext->VSSetShader( m_pVSScene, NULL, 0 );
	pd3dContext->PSSetShader( m_pPSScene, NULL, 0 );
    
    for(UINT i=0; i<pMesh->GetSubsets().size(); i++)
	{
		// Set material value for each subset
		D3D11_MAPPED_SUBRESOURCE MappedResource;
		hr = pd3dContext->Map( m_pcbMatObj, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) ;
		if (FAILED(hr))
			::MessageBoxA(0, "Fail to map constant buffer material", "d3d error", 0);

		CB_MAT_OBJ* pcbPerMaterial = ( CB_MAT_OBJ* )MappedResource.pData;

        const CDX11Mesh::Material* pMat = &(pMesh->GetMaterials()[pMesh->GetSubsets()[i].materialID]);
		pcbPerMaterial->m_MatAmbient = D3DXVECTOR4(pMat->vAmbient, 1);
		pcbPerMaterial->m_MatDiffuse = D3DXVECTOR4(pMat->vDiffuse, -1);
		pcbPerMaterial->m_MatSpecular = D3DXVECTOR4(pMat->vSpecular, float(pMat->nShininess));

		// Set shader resource view
		if(	pMat->pTextureRV11)
		{
			pcbPerMaterial->m_MatDiffuse.w = 1;
			pd3dContext->PSSetShaderResources(IDSRV_TXOBJ, 1, &pMat->pTextureRV11);
		}

		pd3dContext->Unmap( m_pcbMatObj, 0 );

		pd3dContext->VSSetConstantBuffers( 1, 1, &m_pcbMatObj );
		pd3dContext->PSSetConstantBuffers( 1, 1, &m_pcbMatObj );

        int indexCount = pMesh->GetSubsets()[i].triangleCount*3;
        int startLocation = pMesh->GetSubsets()[i].startIndex;
		pd3dContext->DrawIndexed( indexCount, startLocation, 0 );
	}
}

void CSceneRenderer::RenderSceneMeshes(ID3D11DeviceContext* pd3dContext, vector<CDX11Mesh*>& sceneMeshes)
{
    HRESULT hr;

	pd3dContext->IASetInputLayout( m_pLayoutScene );
	pd3dContext->IASetPrimitiveTopology( D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST );
	// set shader
	pd3dContext->VSSetShader( m_pVSScene, NULL, 0 );
	pd3dContext->PSSetShader( m_pPSScene, NULL, 0 );

    // go through all scene meshes
    for (UINT t = 0; t < sceneMeshes.size(); t++)
    {
        CDX11Mesh* pMesh = (sceneMeshes[t]);
        ID3D11Buffer* pVB = pMesh->GetVertexBuffer();
        UINT strides = sizeof( CDX11Mesh::Vertex );
	    UINT offsets = 0;
        pd3dContext->IASetVertexBuffers( 0, 1, &(pVB), &strides, &offsets );
        pd3dContext->IASetIndexBuffer( pMesh->GetIndexBuffer(), DXGI_FORMAT_R32_UINT, 0 );
    
        // go through all subset of mesh i
        for(UINT i=0; i<pMesh->GetSubsets().size(); i++)
	    {
		    // Set material value for each subset
		    D3D11_MAPPED_SUBRESOURCE MappedResource;
		    hr = pd3dContext->Map( m_pcbMatObj, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) ;
		    if (FAILED(hr))
			    ::MessageBoxA(0, "Fail to map constant buffer material", "d3d error", 0);

		    CB_MAT_OBJ* pcbPerMaterial = ( CB_MAT_OBJ* )MappedResource.pData;

            const CDX11Mesh::Material* pMat = &(pMesh->GetMaterials()[pMesh->GetSubsets()[i].materialID]);
		    pcbPerMaterial->m_MatAmbient = D3DXVECTOR4(pMat->vAmbient, 1);
		    pcbPerMaterial->m_MatDiffuse = D3DXVECTOR4(pMat->vDiffuse, -1);
		    pcbPerMaterial->m_MatSpecular = D3DXVECTOR4(pMat->vSpecular, float(pMat->nShininess));

		    // Set shader resource view
		    if(	pMat->pTextureRV11)
		    {
			    pcbPerMaterial->m_MatDiffuse.w = 1;
			    pd3dContext->PSSetShaderResources(IDSRV_TXOBJ, 1, &pMat->pTextureRV11);
		    }

		    pd3dContext->Unmap( m_pcbMatObj, 0 );
		    pd3dContext->PSSetConstantBuffers( 1, 1, &m_pcbMatObj );

            int indexCount = pMesh->GetSubsets()[i].triangleCount*3;
            int startLocation = pMesh->GetSubsets()[i].startIndex;
		    pd3dContext->DrawIndexed( indexCount, startLocation, 0 );
	    }
    }
}

void CSceneRenderer::OnFrameRender(ID3D11DeviceContext* pd3dContext, FRAME_PARAMS& frameParams)
{
    SetSamplerStates(pd3dContext);
    SetConstantBuffer(pd3dContext, frameParams);
    
    RenderSceneMeshes(pd3dContext, m_SceneObjects);
}

void CSceneRenderer::OnDestroy()
{
    // release vertex index buffer etc in each mesh
    for (unsigned int i = 0; i < m_SceneObjects.size(); i++)
        SAFE_DELETE(m_SceneObjects[i]);

    vector<CDX11Mesh*>().swap(m_SceneObjects);
	
    // shader
    SAFE_RELEASE(m_pVSScene);
    SAFE_RELEASE(m_pPSScene);
    // lay out
    SAFE_RELEASE(m_pLayoutScene);
    // constant buffer
    SAFE_RELEASE(m_pcbMatObj);
    SAFE_RELEASE(m_pcbPerFrame);

    // sampler states
	SAFE_RELEASE( m_pSamplerStateLinearWrap );
	SAFE_RELEASE( m_pSamplerStateLinearClamp );
	SAFE_RELEASE( m_pSamplerStatePointClamp );
	SAFE_RELEASE( m_pSamplerStateCmpLess );
}